package com.crazyteam.datasource.annotation;

import com.crazyteam.datasource.dynamic.TransactionExecutor;
import com.crazyteam.datasource.factory.DataSourceFactory;
import com.crazyteam.datasource.dynamic.DynamicDataSource;
import com.crazyteam.datasource.dynamic.DynamicDataSourceTransactionManager;
import com.crazyteam.datasource.config.DataSourceConfigHolder;
import org.apache.commons.lang.StringUtils;
import org.mybatis.spring.SqlSessionFactoryBean;
import org.mybatis.spring.mapper.MapperScannerConfigurer;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.BeansException;
import org.springframework.beans.factory.BeanClassLoaderAware;
import org.springframework.beans.factory.config.BeanDefinition;
import org.springframework.beans.factory.support.BeanDefinitionBuilder;
import org.springframework.beans.factory.support.BeanDefinitionRegistry;
import org.springframework.beans.factory.support.BeanNameGenerator;
import org.springframework.boot.bind.PropertySourceUtils;
import org.springframework.boot.bind.RelaxedPropertyResolver;
import org.springframework.context.*;
import org.springframework.context.annotation.ImportBeanDefinitionRegistrar;
import org.springframework.core.annotation.AnnotationAttributes;
import org.springframework.core.env.Environment;
import org.springframework.core.env.StandardEnvironment;
import org.springframework.core.io.ResourceLoader;
import org.springframework.core.type.AnnotationMetadata;
import org.springframework.jdbc.datasource.DataSourceTransactionManager;
import org.springframework.transaction.PlatformTransactionManager;
import org.springframework.util.ClassUtils;
import org.springframework.util.CollectionUtils;

import javax.sql.DataSource;
import java.beans.Introspector;
import java.util.*;
import java.util.stream.Collectors;
import java.util.stream.Stream;

/**
 * @author guoyang
 * @Date 2018-01-09
 * @Desc 数据源自动装配注册
 */
public class DataSourceImportBeanDefinitionRegistrar implements ImportBeanDefinitionRegistrar,
         BeanClassLoaderAware, EnvironmentAware, ResourceLoaderAware{

    private static final Logger log = LoggerFactory.getLogger(DataSourceImportBeanDefinitionRegistrar.class);

    private static final String DEFAULT_MAPPER_CONFIG = "classpath:ctdatasource/mybatis-configration.xml";

    private ClassLoader classLoader;
    private Environment environment;
    private ResourceLoader resourceLoader;
    private static volatile Boolean hasPrimary = false;
    private static final byte[] primaryLock = new byte[0];

    @Override
    public void setResourceLoader(ResourceLoader resourceLoader) throws BeansException {
        this.resourceLoader = resourceLoader;
    }

    @Override
    public void setBeanClassLoader(ClassLoader classLoader) {
        this.classLoader = classLoader;
    }

    @Override
    public void setEnvironment(Environment environment) {
        this.environment = environment;
    }

    @Override
    public void registerBeanDefinitions(AnnotationMetadata metadata, BeanDefinitionRegistry registry) {
        Map<String, Object> configMetadataAttrs = metadata.getAnnotationAttributes(DataSourceConfig.class.getName(), true);
        AnnotationAttributes configAnnotationAttributes = AnnotationAttributes.fromMap(configMetadataAttrs);
        // 判断是是单个配置还是多个注解
        if (configMetadataAttrs == null || StringUtils.isEmpty(configAnnotationAttributes.getString(DataSourceConfig.ATTR_VALUE))){
            // 多个注解则轮询每个数据源配置注解
            Map<String, Object>  muiltMetadataAttrs = metadata.getAnnotationAttributes(MultiDateSourceConfig.class.getName(), true);
            if (muiltMetadataAttrs == null || muiltMetadataAttrs.isEmpty()){
                throw new RuntimeException("EnableMuiltDateSource 必须配置数据源");
            }else {// [Lorg.springframework.core.annotation.AnnotationAttributes; cannot be cast to [Ljava.lang.annotation.Annotation;
                AnnotationAttributes[] muiltConfigAnnotationAttributes = ( AnnotationAttributes[])muiltMetadataAttrs.get(DataSourceConfig.ATTR_VALUE);
                boolean hasAnyPrimary = false;
                for (AnnotationAttributes configAnnotationAttribute : muiltConfigAnnotationAttributes) {
                    Boolean isPrimary = configAnnotationAttribute.getBoolean(DataSourceConfig.ATTR_IS_PRIMARY);
//                    hasAnyPrimary = hasAnyPrimary || isPrimary;
                    if (isPrimary){
                        hasAnyPrimary = true;
                        break;
                    }
                }
                for (AnnotationAttributes configAnnotationAttribute : muiltConfigAnnotationAttributes) {
                    // 必须配置且只能配置一个 isPrimary 配置项
                    Boolean isPrimary = false;
                    if (!hasPrimary && !hasAnyPrimary){
                        synchronized (primaryLock){
                            if (!hasPrimary && !hasAnyPrimary){
                                isPrimary = true;
                                hasPrimary = true;
                            }
                        }
                    }else {
                        isPrimary = configAnnotationAttribute.getBoolean(DataSourceConfig.ATTR_IS_PRIMARY);
                        if (!hasPrimary && isPrimary){
                            synchronized (primaryLock){
                                if (!hasPrimary && isPrimary){
                                    hasPrimary = true;
                                    isPrimary = true;
                                }
                            }
                        }
                    }
                    registerBeanDefinitionsEach(configAnnotationAttribute, registry, isPrimary);
                }
            }
        }else {
            // 必须配置且只能配置一个 isPrimary 配置项
            Boolean isPrimary = hasPrimary;
            if (!hasPrimary){
                synchronized (primaryLock){
                    if (!hasPrimary ){
                        hasPrimary = true;
                        isPrimary = hasPrimary;
                    }
                }
            }
            registerBeanDefinitionsEach(configAnnotationAttributes, registry, isPrimary);
        }
    }

    private void registerBeanDefinitionsEach(AnnotationAttributes annotationAttributes, BeanDefinitionRegistry registry, Boolean isPrimary) {

        // 读取注解参数
        String prefix = getString(annotationAttributes, DataSourceConfig.ATTR_PREFIX);
        String masterName = getString(annotationAttributes, DataSourceConfig.ATTR_VALUE);
        String[] aliasNames = getStringArray(annotationAttributes, DataSourceConfig.ATTR_ALIAS);
        String[] slaves = getStringArray(annotationAttributes, DataSourceConfig.ATTR_SLAVES);
        String[] mapperLocations = getStringArray(annotationAttributes, DataSourceConfig.ATTR_MAPPER_LOCATIONS);
        String typeAliasesPackage = annotationAttributes.getString(DataSourceConfig.ATTR_TYPE_ALIASES_PACKAGE);
        String configLocation = annotationAttributes.getString(DataSourceConfig.ATTR_CONFIG_LOCATION);
        String[] scannerBasePackages = annotationAttributes.getStringArray(DataSourceConfig.ATTR_SCANNER_BASE_PACKAGES);
        boolean smartSwitch = annotationAttributes.getBoolean(DataSourceConfig.ATTR_SMART_SWITCH);

        String smartSwitchProp = environment.getProperty("datasource.smartSwitch");
        if (StringUtils.isNotBlank(smartSwitchProp)){
            smartSwitch = "true".equalsIgnoreCase(smartSwitchProp);
        }

        Class<? extends DataSource> datasourceType = getToClass(annotationAttributes, DataSourceConfig.ATTR_DATASOURCE_TYPE, DataSource.class);

        DataSourceConfigHolder dataSourceConfig = new DataSourceConfigHolder(masterName, slaves, aliasNames);

        // 注册主库数据源
        registDataSource(prefix, masterName, datasourceType, registry, isPrimary);
        // 注册主库 Master别名
        registry.registerAlias(dataSourceConfig.getDataSourceBeanName(), dataSourceConfig.getMasterDataSourceBeanName());

        // 注册主库事务
        registTransactionManager(dataSourceConfig.getTransactionManagerBeanName(), dataSourceConfig.getDataSourceBeanName(), null, registry, isPrimary);

        // 注册主库事务执行器
        registTransactionExecutor(dataSourceConfig.getTransactionExecutorBeanName(), dataSourceConfig.getTransactionManagerBeanName(), registry, isPrimary);

        // 注册主库SqlSession
        registSqlSession(dataSourceConfig.getSqlSessionFactoryBeanName(), dataSourceConfig.getDataSourceBeanName(), typeAliasesPackage, mapperLocations, configLocation, registry, isPrimary);

        // 注册主库扫描器
        registMapperScannerConfigurer(dataSourceConfig.getMapperScannerBeanName(), dataSourceConfig.getSqlSessionFactoryBeanName(), scannerBasePackages, registry, new BeanNameGenerator() {
            @Override
            public String generateBeanName(BeanDefinition definition, BeanDefinitionRegistry beanDefinitionRegistry) {
                String shortClassName = ClassUtils.getShortName(definition.getBeanClassName());
                // 主库 mapper 设为 Primary
                definition.setPrimary(true);
                String mapperName = Introspector.decapitalize(shortClassName);
                // 注册主库 mapper别名
                registry.registerAlias(mapperName, "master" + shortClassName);
                return mapperName;
            }
        });
        // 注册主库别名
        dataSourceConfig.getAliasConfig().forEach((aliasConf) -> {
            registry.registerAlias(dataSourceConfig.getDataSourceBeanName(), aliasConf.getDataSourceBeanName());
            registry.registerAlias(dataSourceConfig.getMasterDataSourceBeanName(), aliasConf.getMasterDataSourceBeanName());
//            registry.registerAlias(dataSourceConfig.getSqlSessionFactoryBeanName(), aliasConf.getSqlSessionFactoryBeanName());
            registry.registerAlias(dataSourceConfig.getTransactionManagerBeanName(), aliasConf.getTransactionManagerBeanName());
            registry.registerAlias(dataSourceConfig.getTransactionExecutorBeanName(), aliasConf.getTransactionExecutorBeanName());
//            registry.registerAlias(dataSourceConfig.getMapperScannerBeanName(), aliasConf.getMapperScannerBeanName());
        });

        // 有备库
        if (!CollectionUtils.isEmpty(dataSourceConfig.getSlavesDataSourceBeanNames())) {
            // 注册备库
            dataSourceConfig.getSlavesDataSourceNames().forEach((slave) -> registDataSource(prefix, slave, datasourceType, registry, false));

            // 注册动态备库
            registDynamicSlaveDataSource(dataSourceConfig, registry, smartSwitch);

            // 注册动态备库事务
            registTransactionManager(dataSourceConfig.getSlaveTransactionManagerBeanName(), dataSourceConfig.getDynamicSlaveDataSourceBeanName(), DynamicDataSourceTransactionManager.class, registry, false);
            // 注册动态备库事务执行器
            registTransactionExecutor(dataSourceConfig.getSlaveTransactionExecutorBeanName(), dataSourceConfig.getSlaveTransactionManagerBeanName(), registry, isPrimary);

            // 动态数据库必须配置mybatise配置项,如未配置，则默认
            if (StringUtils.isEmpty(configLocation)){
                configLocation = DEFAULT_MAPPER_CONFIG;
            }
            // 注册动态备库 SqlSession
            registSqlSession(dataSourceConfig.getSlaveSqlSessionFactoryBeanName(), dataSourceConfig.getDynamicSlaveDataSourceBeanName(), typeAliasesPackage, mapperLocations, configLocation, registry, false);
            // 注册动态备库扫描器
            registMapperScannerConfigurer(dataSourceConfig.getSlaveMapperScannerBeanName(), dataSourceConfig.getSlaveSqlSessionFactoryBeanName(), scannerBasePackages, registry, new BeanNameGenerator() {
                @Override
                public String generateBeanName(BeanDefinition definition, BeanDefinitionRegistry beanDefinitionRegistry) {
                    String shortClassName = ClassUtils.getShortName(definition.getBeanClassName());
                    // 主库未primary，备库全部未false
                    definition.setPrimary(false);
                    return "slave" + shortClassName;
                }
            });

            // 注册主库别名
            dataSourceConfig.getAliasConfig().forEach((aliasConf) -> {
                registry.registerAlias(dataSourceConfig.getDynamicSlaveDataSourceBeanName(), aliasConf.getDynamicSlaveDataSourceBeanName());
                registry.registerAlias(dataSourceConfig.getSlaveSqlSessionFactoryBeanName(), aliasConf.getSlaveSqlSessionFactoryBeanName());
                registry.registerAlias(dataSourceConfig.getSlaveTransactionManagerBeanName(), aliasConf.getSlaveTransactionManagerBeanName());
                registry.registerAlias(dataSourceConfig.getSlaveTransactionExecutorBeanName(), aliasConf.getSlaveTransactionExecutorBeanName());
                registry.registerAlias(dataSourceConfig.getSlaveMapperScannerBeanName(), aliasConf.getSlaveMapperScannerBeanName());
            });

        }

    }

    private BeanDefinition registMapperScannerConfigurer(String beanName, String sqlSessionFactoryBeanName, String[] scannerBasePackages, BeanDefinitionRegistry registry, BeanNameGenerator nameGenerator) {
        if (isEmptyStrArr(scannerBasePackages)){
            return null;
        }
        BeanDefinitionBuilder mapperScannerDefinitionBuilder = BeanDefinitionBuilder.genericBeanDefinition(MapperScannerConfigurer.class);
        mapperScannerDefinitionBuilder.addPropertyValue("basePackage", StringUtils.join(scannerBasePackages, ConfigurableApplicationContext.CONFIG_LOCATION_DELIMITERS));
        mapperScannerDefinitionBuilder.addPropertyValue("sqlSessionFactoryBeanName", sqlSessionFactoryBeanName);
        mapperScannerDefinitionBuilder.addPropertyValue("nameGenerator", nameGenerator);
        mapperScannerDefinitionBuilder.addDependsOn(sqlSessionFactoryBeanName);
        BeanDefinition beanDefinition = mapperScannerDefinitionBuilder.getBeanDefinition();
        beanDefinition.setLazyInit(false);
        registry.registerBeanDefinition( beanName , beanDefinition);

        return beanDefinition;
    }

    private BeanDefinition registSqlSession(String sqlSessionFactoryBeanName, String dataSourceBeanName, String typeAliasesPackage, String[] mapperLocations, String configLocation, BeanDefinitionRegistry registry, Boolean isPrimary) {
        if (isEmptyStrArr(mapperLocations)){
            return null;
        }
        BeanDefinitionBuilder sqlSessionDefinitionBuilder = BeanDefinitionBuilder.genericBeanDefinition(SqlSessionFactoryBean.class);

        sqlSessionDefinitionBuilder.addPropertyReference("dataSource", dataSourceBeanName);
        sqlSessionDefinitionBuilder.addPropertyValue("typeAliasesPackage", typeAliasesPackage);

        // TODO
        if (StringUtils.isNotEmpty(configLocation)){
            sqlSessionDefinitionBuilder.addPropertyValue("configLocation", resourceLoader.getResource(configLocation));
        }
        sqlSessionDefinitionBuilder.addPropertyValue("mapperLocations", mapperLocations);

        sqlSessionDefinitionBuilder.addDependsOn(dataSourceBeanName);

        BeanDefinition beanDefinition = sqlSessionDefinitionBuilder.getBeanDefinition();
        beanDefinition.setPrimary(isPrimary);
        beanDefinition.setLazyInit(false);
        registry.registerBeanDefinition( sqlSessionFactoryBeanName , beanDefinition);

        return beanDefinition;
    }

    private BeanDefinition registDynamicSlaveDataSource(DataSourceConfigHolder dataSourceConfig, BeanDefinitionRegistry registry, Boolean smartSwitch) {
        BeanDefinitionBuilder routeDataSourceDefinitionBuilder = BeanDefinitionBuilder.genericBeanDefinition(DynamicDataSource.class);
        routeDataSourceDefinitionBuilder.addConstructorArgValue(dataSourceConfig.getMasterDataSourceBeanName());
        routeDataSourceDefinitionBuilder.addConstructorArgValue(dataSourceConfig.getSlavesDataSourceBeanNames());
        routeDataSourceDefinitionBuilder.addPropertyValue("smartSwitch", smartSwitch);
        routeDataSourceDefinitionBuilder.addDependsOn(dataSourceConfig.getMasterDataSourceBeanName());
        dataSourceConfig.getSlavesDataSourceBeanNames().forEach(routeDataSourceDefinitionBuilder::addDependsOn);
        BeanDefinition beanDefinition = routeDataSourceDefinitionBuilder.getBeanDefinition();
        beanDefinition.setPrimary(false);
        beanDefinition.setLazyInit(false);
        registry.registerBeanDefinition( dataSourceConfig.getDynamicSlaveDataSourceBeanName() , beanDefinition);

        return beanDefinition;

    }

    private BeanDefinition registTransactionManager(String beanName, String dataSourceBeanName, Class<? extends PlatformTransactionManager> transactionManagerType, BeanDefinitionRegistry registry, Boolean isPrimary) {
        if (transactionManagerType == null){
            transactionManagerType = DataSourceTransactionManager.class;
        }
        BeanDefinitionBuilder transactionDefinitionBuilder = BeanDefinitionBuilder.genericBeanDefinition(transactionManagerType);
        transactionDefinitionBuilder.addPropertyReference("dataSource", dataSourceBeanName);
        transactionDefinitionBuilder.addDependsOn(dataSourceBeanName);

        BeanDefinition beanDefinition = transactionDefinitionBuilder.getBeanDefinition();
        beanDefinition.setPrimary(isPrimary);
        beanDefinition.setLazyInit(false);
        registry.registerBeanDefinition(beanName , beanDefinition);

        return beanDefinition;
    }
    private BeanDefinition registTransactionExecutor(String beanName, String transactionManagerBeanName, BeanDefinitionRegistry registry, Boolean isPrimary) {
        BeanDefinitionBuilder transactionDefinitionBuilder = BeanDefinitionBuilder.genericBeanDefinition(TransactionExecutor.class);
        transactionDefinitionBuilder.addPropertyReference("transactionManager", transactionManagerBeanName);

        BeanDefinition beanDefinition = transactionDefinitionBuilder.getBeanDefinition();
        beanDefinition.setPrimary(isPrimary);
        beanDefinition.setLazyInit(false);
        registry.registerBeanDefinition(beanName , beanDefinition);

        return beanDefinition;

    }

    private BeanDefinition registDataSource(String prefix, String datasourceName, Class<? extends DataSource> datasourceType, BeanDefinitionRegistry registry, Boolean isPrimary) {
        Map<String, Object> dataSourceProperties = new RelaxedPropertyResolver(environment).getSubProperties(prefix + datasourceName + ".");

        PropertySourceUtils.getSubProperties(((StandardEnvironment)environment).getPropertySources(), prefix + datasourceName + ".");
        dataSourceProperties = new HashMap<>(dataSourceProperties);
        dataSourceProperties.put("name" , datasourceName);

        // 注册数据源lanDefinitionBuilder
        BeanDefinitionBuilder dataSourceDefinitionBuilder = BeanDefinitionBuilder.genericBeanDefinition(DataSourceFactory.class);
        dataSourceDefinitionBuilder.addConstructorArgValue(datasourceType);
        dataSourceDefinitionBuilder.addConstructorArgValue(dataSourceProperties);

        BeanDefinition beanDefinition = dataSourceDefinitionBuilder.getBeanDefinition();
        beanDefinition.setPrimary(isPrimary);
        beanDefinition.setLazyInit(false);
        registry.registerBeanDefinition( DataSourceConfigHolder.buildDataSourceBeanName(datasourceName) , beanDefinition);

        return beanDefinition;
    }

     private String getString(AnnotationAttributes annotationAttributes, String atter){
        return annotationAttributes.getString(atter);
     }

    private String[] getStringArray(AnnotationAttributes annotationAttributes, String atter){
        List<String> resultList = Stream.of(annotationAttributes.getStringArray(atter)).filter((s) -> StringUtils.isNotEmpty(s)).collect(Collectors.toList());
        return resultList.toArray(new String[resultList.size()]);
    }

    private static boolean isEmptyStrArr(String... vals){
        if (vals == null || vals.length == 0){
            return true;
        }else {
            for (String val:vals){
                if (StringUtils.isNotBlank(val)){
                    return false;
                }
            }
            return true;
        }
    }

    private  <T> Class<? extends T> getToClass(AnnotationAttributes annotationAttributes, String atter, Class<T> classType){
        Class<? extends T> sourceType = null;
        String className = getString(annotationAttributes, atter);
        try {
            Class type =  ClassUtils.forName(className, classLoader);
            if (classType.isAssignableFrom(type)){
                sourceType = type;
            }else {
                throw new RuntimeException("不支持的配置" + className);
            }
        } catch (Exception e) {
            log.error("不支持的配置" + className, e);
            throw new RuntimeException("不支持的配置" + className, e);
        }
        return sourceType;
    }
}
